"""
使用PandasAI分析CSV文件
支持自然语言查询数据
"""

import os
import pandas as pd
from pandasai import SmartDataframe
from pandasai.llm import OpenAI
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

def analyze_csv_with_pandasai(csv_path, chunk_size=10000):
    """
    使用PandasAI分析CSV文件
    
    Args:
        csv_path: CSV文件路径
        chunk_size: 分块大小，用于处理大文件
    """
    
    # 检查API密钥
    api_key = os.getenv('OPENAI_API_KEY')
    if not api_key:
        print("错误: 未找到OPENAI_API_KEY，请在.env文件中设置")
        return
    
    # 初始化LLM
    llm = OpenAI(api_token=api_key)
    
    print(f"正在加载CSV文件: {csv_path}")
    print(f"使用分块大小: {chunk_size} 行")
    
    try:
        # 先读取一小部分数据进行分析
        df_sample = pd.read_csv(csv_path, nrows=chunk_size, encoding='utf-8')
        
        print(f"\n数据概览:")
        print(f"- 列数: {len(df_sample.columns)}")
        print(f"- 列名: {list(df_sample.columns)}")
        print(f"\n前5行数据:")
        print(df_sample.head())
        
        # 创建SmartDataframe
        smart_df = SmartDataframe(df_sample, config={"llm": llm})
        
        print("\n" + "="*50)
        print("PandasAI 智能分析系统已启动")
        print("="*50)
        print("\n提示: 你可以用自然语言提问，例如:")
        print("  - 这个数据集有多少行？")
        print("  - 显示前10行数据")
        print("  - 哪一列的数值最大？")
        print("  - 绘制数据分布图")
        print("\n输入 'quit' 或 'exit' 退出\n")
        
        # 交互式查询循环
        while True:
            query = input("请输入你的问题: ").strip()
            
            if query.lower() in ['quit', 'exit', '退出']:
                print("再见！")
                break
            
            if not query:
                continue
            
            try:
                response = smart_df.chat(query)
                print(f"\n回答: {response}\n")
            except Exception as e:
                print(f"查询出错: {e}\n")
    
    except Exception as e:
        print(f"处理CSV文件时出错: {e}")


def main():
    """主函数"""
    
    # 可用的CSV文件
    csv_files = [
        "output/基本设施.csv",
        "output/设施集合.csv"
    ]
    
    print("可用的CSV文件:")
    for i, file in enumerate(csv_files, 1):
        file_size = os.path.getsize(file) / (1024 * 1024)  # 转换为MB
        print(f"{i}. {file} ({file_size:.2f} MB)")
    
    # 选择文件
    try:
        choice = int(input("\n请选择要分析的文件 (输入序号): "))
        if 1 <= choice <= len(csv_files):
            selected_file = csv_files[choice - 1]
            
            # 对于大文件，使用较小的chunk_size
            chunk_size = 10000  # 可以根据需要调整
            
            analyze_csv_with_pandasai(selected_file, chunk_size)
        else:
            print("无效的选择")
    except ValueError:
        print("请输入有效的数字")
    except KeyboardInterrupt:
        print("\n\n程序已终止")


if __name__ == "__main__":
    main()

